perf: in-place ops cut memory-read VRAM - #1

Summary

Applies the VRAM-optimization pattern from "Optimizing PyTorch VRAM usage for mcbyte-tracker (16GB → under 8GB)" to Cutie's memory-read hot path: replacing chained non-in-place tensor arithmetic with in-place and fused ops, which the article found to be the actual VRAM win — not del + torch.cuda.empty_cache(), which that article measured as both slower and less stable than doing nothing.

get_similarity() in cutie/model/utils/memory_utils.py had the exact same anti-pattern as the article's baseline: several full-size B×N×HW intermediate tensors (a_sq, two_ab, b_sq) built via chained -a_sq + two_ab - b_sq, recomputed every frame against the full accumulated memory bank. A follow-up pattern search across cutie/ surfaced five more sites with the same shape, now fused or converted to in-place forms as well.

Test-first throughout: characterization tests (forward-numeric snapshots, backward-gradient snapshots, and CPU-vs-MPS parity) were written and confirmed green against the unmodified source before any production edit landed, per this repo's refactor discipline. See the companion commit test(cutie): guardrail tests before VRAM refactor.

Changes

• cutie/model/utils/memory_utils.py — get_similarity(): fuse the 2 * scalar multiply into two_ab.mul_(2), chain .sub_() for the a_sq/b_sq combination instead of allocating separate -a_sq + two_ab - b_sq intermediates. do_softmax() dense branch: .exp_() on the throwaway sub() output instead of a separately-allocated torch.exp() tensor.
• cutie/model/group_modules.py — MainToGroupDistributor (method='muladd'): x * g + g → torch.addcmul(g, x, g), one fused kernel instead of two allocations.
• cutie/model/channel_attn.py — CAResBlock.forward: x * w + downsample(r) → torch.addcmul(downsample(r), x, w).
• cutie/model/big_modules.py — KeyProjection shrinkage branch: d_proj(x) ** 2 + 1 → d_proj(x).pow(2).add_(1).
• cutie/model/cutie.py — encode_image/encode_mask normalization: (image - mean) / std → image.sub(mean).div_(std). sub() always copies, so the caller's frame tensor is never mutated — only the fresh copy is divided in place.
• cutie/inference/memory_manager.py — streaming object-value accumulation now uses .add_() in place instead of building a new_acc temp and reassigning both memory slices.
• scripts/bench_vram_mps.py (new) — standalone MPS peak-memory and wall-clock bench for the memory-read path. Not a pytest gate: MPS allocator readings are noisy run-to-run, so this is evidence to eyeball, not a CI assertion.

Why not every op went in-place

get_similarity/do_softmax are shared by the training path (cutie.py → get_affinity, needs autograd) and the inference path (memory_manager.py, no grad) — an in-place op that overwrites a tensor autograd needs for backward silently breaks training. One candidate op was reverted from in-place back to out-of-place after the backward-parity guardrail caught it: x_exp.div_(x_exp_sum) in do_softmax's dense branch raised RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation — Exp's own backward needs its output preserved, so that final normalization stays out-of-place (documented inline). The guardrail test, not assumption, decided every in-place/out-of-place call in this PR.

MPS bench (this MacBook, torch 2.13, N=32400 accumulated memory entries)

~49% peak-memory cut, ~5% faster — consistent with the article's finding that structural in-place redesign outperforms aggressive cache clearing.

Test coverage

• tests/test_memory_utils.py — extended with 21 characterization cases: get_similarity (qe present/None, ms present/None, add_batch_dim, CK=1/N=1 boundaries), do_softmax (dense branch, top_k == N boundary, inplace identity, tie-break, large-magnitude stability), readout (T=1/T>1), plus forward and backward-gradient snapshots for every op on the training path.
• tests/test_model_ops_vram.py (new) — forward/backward parity for CAResBlock, MainToGroupDistributor(muladd), KeyProjection(need_s), and a caller-tensor-safety test for the cutie.py normalize idiom.
• tests/test_mps_parity.py (new) — CPU-vs-MPS numeric parity for get_similarity, do_softmax, and CAResBlock, actually executed on Apple GPU hardware (not skipped) — the first MPS coverage in this repo's test suite, which was previously CPU-only end to end.
• Full suite: 88 passed. ruff check cutie/ tests/ scripts/: clean.

Known limitation (not fixed here, out of scope)

do_softmax's top-k branch already breaks .backward() on main, regardless of the inplace flag — an unconditional values.exp_() on the torch.topk output. This is a pre-existing bug, unrelated to this refactor, pinned as-is by test_do_softmax_top_k_branch_backward_raises_inplace_version_error. No live impact today since the top-k path is only exercised inference-side (memory_manager.py, no grad); would need a separate fix if the top-k path is ever used with requires_grad.